import re
from c2rust.actions import *
from metagpt.roles import Role
from metagpt.schema import Message
from metagpt.logs import logger
from metagpt.actions import UserRequirement

def get_class_name(cls) -> str:
    """Return class name"""
    return f"{cls.__module__}.{cls.__name__}"
class C_Commentator(Role):
    name: str = "C_Commentator"
    profile: str = "C_Commentator"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._watch([UserRequirement])
        self.set_actions([CommentCode])
        
    async def _act(self) -> Message:
        logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
        todo = self.rc.todo
        
        context = self.get_memories(k=1)
        # context = self.get_memories()  # use all memories as context
        logger.info(f"context: {context}")
        code_text = await todo.run(TEAM_DESCRIPTION, context[0].content)  # specify arguments
        msg = Message(content=code_text, role=self.profile, cause_by=type(todo))

        return msg
    

class RustTranslator(Role):
    name: str = "RustTranslator"
    profile: str = "RustTranslator"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._watch([CommentCode])
        self.set_actions([TranslateCode])
        
    async def _act(self) -> Message:
        logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
        todo = self.rc.todo

        context = self.get_memories(k=1)
        # context = self.get_memories()  # use all memories as context
        logger.info(f"context: {context}")
        code_text = await todo.run(TEAM_DESCRIPTION, context[0].content)  # specify arguments
        msg = Message(content=code_text, role=self.profile, cause_by=type(todo))

        return msg
    
class AddMainFunction(Role):
    name: str = "AddMainFunction"
    profile: str = "AddMainFunction"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._watch([TranslateCode])
        self.set_actions([AddMainFunc])

    async def _act(self) -> Message:
        logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
        todo = self.rc.todo

        context = self.get_memories(k=1)
        # context = self.get_memories()  # use all memories as context
        logger.info(f"context: {context}")
        code_text = await todo.run(TEAM_DESCRIPTION, context[0].content)  # specify arguments
        msg = Message(content=code_text, role=self.profile, cause_by=type(todo))

        return msg

class RustEnhancer(Role):
    name: str = "RustEnhancer"
    profile: str = "RustEnhancer"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._watch([AddMainFunc, CorrectCode])
        self.set_actions([EnhanceCode, EnhanceCodeAfterCorrect])
        
    async def _think(self):
        news = self.rc.news[0]
        logger.info(news.cause_by)
        logger.info(type(news.cause_by))
        logger.info(get_class_name(CorrectCode))
        if news.cause_by == get_class_name(CorrectCode):
            self.rc.todo = EnhanceCodeAfterCorrect()
        else:
            self.rc.todo = EnhanceCode()
        logger.info(self.rc.todo)

    async def _act(self) -> Message:
        logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
        todo = self.rc.todo

        context = self.get_memories(k=1)
        # context = self.get_memories()  # use all memories as context
        logger.info(f"context: {context}")
        code_text = await todo.run(TEAM_DESCRIPTION, context[0].content)  # specify arguments
        msg = Message(content=code_text, role=self.profile, cause_by=type(todo))

        return msg
        
class RustRunner(Role):
    name: str = "RustRunner"
    profile: str = "RustRunner"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._watch([EnhanceCode, EnhanceCodeAfterCorrect, EnhanceCodeRAG, EnhanceCodeAfterCorrectRAG])
        self.set_actions([RunCode])
        self.has_dep = False

    async def _act(self) -> Message:
        logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
        todo = self.rc.todo

        context = self.get_memories(k=1)
        # context = self.get_memories()  # use all memories as context
        logger.info(f"context: {context}")
        file_name = await todo.run(context[0].content)  # specify arguments
        msg = Message(content=file_name, role=self.profile, cause_by=type(todo))

        return msg
    
class CompilerMessageAnalyser(Role):
    name: str = "CompilerMessageAnalyser"
    profile: str = "CompilerMessageAnalyser"

    def __init__(self, epoch_limit=5, **kwargs):
        super().__init__(**kwargs)
        self.epoch = 0
        self.epoch_limit = epoch_limit
        self._watch([RunCode])
        self.set_actions([AnnounceError, AnnounceSuccess, AnnounceEnd])
        
    async def _think(self):
        news = self.rc.news[0]
        self.epoch += 1
        filenames = news.content
        filenames = filenames.strip()
        result_file = filenames.split(";")[-1]
        with open(result_file, "r") as f:
            compiler_message = f.read()
        
        pattern = r"error"
        match = re.search(pattern, compiler_message, re.DOTALL)
        if not match:
            logger.info("Success after {} epoches".format(self.epoch))
            with open("epoch.txt", "w") as f1:
                f1.write(str(self.epoch))
            self.rc.todo = AnnounceSuccess()
            self.epoch = 0
        else:
            self.rc.todo = AnnounceError()
            if self.epoch >= self.epoch_limit:
                with open("epoch.txt", "w") as f1:
                    f1.write(str(self.epoch))
                self.rc.todo = AnnounceEnd()
                self.epoch = 0

    async def _act(self) -> Message:
        logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
        todo = self.rc.todo
        
        if isinstance(todo, AnnounceSuccess):
            return Message(content="Congratulations! There is no error in compiler's messages.", role=self.profile, cause_by=type(todo))
        elif isinstance(todo, AnnounceEnd):
            return Message(content="Reached epoch limit", role=self.profile, cause_by=type(todo))
        else:
            context = self.get_memories(k=1)
            # context = self.get_memories()  # use all memories as context
            logger.info(f"context: {context}")
            code_text = await todo.run(context[0].content)  # specify arguments
            msg = Message(content=code_text, role=self.profile, cause_by=type(todo))
            return msg
        
class RustCorrector(Role):
    name: str = "RustCorrector"
    profile: str = "RustCorrector"

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self._watch([AnnounceError])
        self.set_actions([CorrectCode])

    async def _act(self) -> Message:
        logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
        todo = self.rc.todo

        context = self.get_memories(k=1)
        # context = self.get_memories()  # use all memories as context
        logger.info(f"context: {context}")
        code_text = await todo.run(TEAM_DESCRIPTION, context[0].content)  # specify arguments
        msg = Message(content=code_text, role=self.profile, cause_by=type(todo))

        return msg
    
# class TestModifier(Role):
#     name: str = "TestModifier"
#     profile: str = "TestModifier"

#     def __init__(self, **kwargs):
#         super().__init__(**kwargs)
#         self._watch([AnnounceSuccess, AnnounceEnd])
#         self.set_actions([ModifyTest])

#     async def _act(self) -> Message:
#         logger.info(f"{self._setting}: to do {self.rc.todo}({self.rc.todo.name})")
#         todo = self.rc.todo

#         context = self.get_memories(k=1)
#         # context = self.get_memories()  # use all memories as context
#         logger.info(f"context: {context}")
#         code_text = await todo.run(TEAM_DESCRIPTION, context[0].content)  # specify arguments
#         msg = Message(content=code_text, role=self.profile, cause_by=type(todo))

#         return msg
